import json
import os
from urllib.request import urlopen
import torch
from torch.utils.data import Dataset
import tiktoken
from functools import partial
from torch.utils.data import DataLoader
from gpt2_download import download_and_load_gpt2
from gpt2_tools import GPTModel, load_weights_into_gpt, generate, text_to_token_ids, token_ids_to_text, calc_loss_loader, train_model_simple
import time
import re

def download_and_load_file(file_path, url):
    if not os.path.exists(file_path):
        with urlopen(url) as response:
            text_data = response.read().decode("utf-8")
            with open(file_path, "w", encoding="utf-8") as file:
                file.write(text_data)
    else:
        with open(file_path, "r", encoding="utf-8") as file:
            text_data = file.read()
    with open(file_path, "r") as file:
        data = json.load(file)
    return data

file_path = "instruction-data.json"
url = (
    "https://raw.githubusercontent.com/rasbt/LLMs-from-scratch"
    "/main/ch07/01_main-chapter-code/instruction-data.json"
)

data = download_and_load_file(file_path, url)
# print("Number of entries:", len(data))
#
# '''
# Number of entries: 1100
# '''
#
# print("Example entry:\n", data[50])
# '''
# Example entry:
#  {'instruction': 'Identify the correct spelling of the following word.',
#   'input': 'Ocassion', 'output': "The correct spelling is 'Occasion.'"}
# '''
#
# print("Another example entry:\n", data[999])
# # 根据该条目的内容，input 字段有时可能为空：
# '''
# {'instruction': 'What is an antonym of "complicated"?',
#  'input': '',
#  'output': 'An antonym of "complicated" is "simple."'}
# '''

def format_input(entry):
    instruction_text = (
        f"Below is an instruction that describes a task. "
        f"Write a response that appropriately completes the request."
        f"\n\n### Instruction:\n{entry['instruction']}"
    )

    input_text = (
        f"\n\n### Input:\n{entry['input']}" if entry["input"] else ""
    )

    return instruction_text + input_text

# model_input = format_input(data[50])
# desired_response = f"\n\n### Response:\n{data[50]['output']}"
# print(model_input + desired_response)
# '''
# Below is an instruction that describes a task. Write a response that
# appropriately completes the request.
#  ### Instruction:
#  Identify the correct spelling of the following word.
#  ### Input:
#  Ocassion
#  ### Response:
#  The correct spelling is 'Occasion.'
# '''
#
# model_input = format_input(data[999])
# desired_response = f"\n\n### Response:\n{data[999]['output']}"
# print(model_input + desired_response)
# # 输出显示，input 字段为空的条目在格式化的输入中不会包含 ### Input: 部分。
# '''
# Below is an instruction that describes a task. Write a response that
# appropriately completes the request.
# ### Instruction:
# What is an antonym of 'complicated'?
# ### Response:
# An antonym of 'complicated' is 'simple'.
# '''

train_portion = int(len(data) * 0.85)
test_portion = int(len(data) * 0.1)
val_portion = len(data) - train_portion - test_portion

train_data = data[:train_portion]
test_data = data[train_portion:train_portion + test_portion]
val_data = data[train_portion + test_portion:]

# print("Training set length:", len(train_data))
# print("Validation set length:", len(val_data))
# print("Test set length:", len(test_data))
# '''
#  Training set length: 935
#  Validation set length: 55
#  Test set length: 110
# '''

class InstructionDataset(Dataset):
    def __init__(self, data, tokenizer):
        self.data = data
        self.encoded_texts = []  # 用于存储编码后的文本
        for entry in data:
            # 格式化输入条目
            instruction_plus_input = format_input(entry)
            # 格式化响应文本
            response_text = f"\n\n### Response:\n{entry['output']}"
            # 将指令和响应合并为完整文本
            full_text = instruction_plus_input + response_text
            # 对完整文本进行分词并编码，添加到编码后的文本列表中
            self.encoded_texts.append(
                tokenizer.encode(full_text)
            )

    def __getitem__(self, index):
        # 获取指定索引的编码文本
        return self.encoded_texts[index]

    def __len__(self):
        # 返回数据集的长度
        return len(self.data)

tokenizer = tiktoken.get_encoding("gpt2")



# inputs_1 = [0, 1, 2, 3, 4]
# inputs_2 = [5, 6]
# inputs_3 = [7, 8, 9]
# batch = (
#    inputs_1,
#    inputs_2,
#    inputs_3
# )

def custom_collate_fn(
    batch,
    pad_token_id=50256,  # 填充标记的ID
    ignore_index=-100,   # 忽略的索引，用于损失计算
    allowed_max_length=None,  # 允许的最大序列长度
    device="cpu"         # 目标设备，默认是CPU
):
    # 找到批次中最长的序列长度并加1
    batch_max_length = max(len(item) + 1 for item in batch)
    inputs_lst, targets_lst = [], []  # 初始化输入列表和目标列表

    # 遍历批次中的每个序列，进行填充处理
    for item in batch:
        new_item = item.copy()  # 复制当前序列
        new_item += [pad_token_id]  # 在序列末尾添加一个填充标记

        # 计算剩余需要填充的长度
        padded = (
            new_item + [pad_token_id] *
            (batch_max_length - len(new_item))
        )

        # 将填充后的序列转换为tensor
        inputs = torch.tensor(padded[:-1])  # 去掉最后一个填充标记，作为输入
        targets = torch.tensor(padded[1:])  # 向右偏移一个标记，作为目标

        # 找到填充标记的索引，并将除第一个填充标记外的目标标记替换为忽略索引
        mask = targets == pad_token_id
        indices = torch.nonzero(mask).squeeze()
        if indices.numel() > 1:
            targets[indices[1:]] = ignore_index

        # 如果设置了允许的最大长度，则对输入和目标进行裁剪
        if allowed_max_length is not None:
            inputs = inputs[:allowed_max_length]
            targets = targets[:allowed_max_length]

        # 将处理后的输入和目标tensor分别添加到对应列表中
        inputs_lst.append(inputs)
        targets_lst.append(targets)

    # 将输入和目标列表中的所有tensor堆叠为一个tensor，并传输到目标设备
    inputs_tensor = torch.stack(inputs_lst).to(device)
    targets_tensor = torch.stack(targets_lst).to(device)
    return inputs_tensor, targets_tensor  # 返回输入和目标的tensor

# 使用自定义的collate函数对批次进行处理
# inputs, targets = custom_collate_fn(batch)
# print(inputs)  # 打印输入tensor
# print(targets)  # 打印目标tensor
#
# '''
# tensor([[    0,     1,     2,    3,     4],
#          [    5,     6, 50256, 50256, 50256],
#          [    7,     8,     9, 50256, 50256]])
#
# tensor([[    1,     2,     3,    4, 50256],
#         [    6, 50256,  -100,  -100,  -100],
#         [    8,     9, 50256,  -100,  -100]])
# '''

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
customized_collate_fn = partial(
    custom_collate_fn,
    device=device,
    allowed_max_length=1024
)

num_workers = 0
batch_size = 8

torch.manual_seed(123)

train_dataset = InstructionDataset(train_data, tokenizer)
train_loader = DataLoader(
    train_dataset,
    batch_size=batch_size,
    collate_fn=customized_collate_fn,
    shuffle=True,
    drop_last=True,
    num_workers=num_workers
)

val_dataset = InstructionDataset(val_data, tokenizer)
val_loader = DataLoader(
    val_dataset,
    batch_size=batch_size,
    collate_fn=customized_collate_fn,
    shuffle=False,
    drop_last=False,
    num_workers=num_workers
)

test_dataset = InstructionDataset(test_data, tokenizer)
test_loader = DataLoader(
    test_dataset,
    batch_size=batch_size,
    collate_fn=customized_collate_fn,
    shuffle=False,
    drop_last=False,
    num_workers=num_workers
)
#
# print("Train loader:")
# for inputs, targets in train_loader:
#     print(inputs.shape, targets.shape)

BASE_CONFIG={
    "vocab_size": 50257, "drop_rate": 0.0, "qkv_bias": True, "context_length": 1024
}

model_configs={
    "gpt2-small (124M)": {"emb_dim": 768, "n_layers": 12, "n_heads":12},
    "gpt2-medium (355M)": {"emb_dim": 1024, "n_layers": 24, "n_heads": 16},
    "gpt2-large (774M)": {"emb_dim": 1280, "n_layers": 36, "n_heads": 20},
    "gpt2-x1 (1558M)":{"emb_dim": 1600, "n_layers":48, "n_heads": 25}
}
CHOOSE_MODEL = "gpt2-medium (355M)"
BASE_CONFIG.update(model_configs[CHOOSE_MODEL])

model_size = CHOOSE_MODEL.split(" ")[-1].lstrip("(").rstrip(")")
settings, params = download_and_load_gpt2(
    model_size=model_size,
    models_dir="gpt2"
)
model = GPTModel(BASE_CONFIG)
load_weights_into_gpt(model, params)
model.eval()

torch.manual_seed(123)
input_text = format_input(val_data[0])
print(input_text)

token_ids = generate(
    model=model,
    idx=text_to_token_ids(input_text, tokenizer),
    max_new_tokens=35,
    context_size=BASE_CONFIG["context_length"],
    eos_id=50256,
)
generated_text = token_ids_to_text(token_ids, tokenizer)

response_text = generated_text[len(input_text):].strip()
print(response_text)

model.to(device)
torch.manual_seed(123)
with torch.no_grad():
    train_loss = calc_loss_loader(
        train_loader, model, device, num_batches=5
    )
    val_loss = calc_loss_loader(
        val_loader, model, device, num_batches=5
    )
print("Training loss:", train_loss)
print("Validation loss:", val_loss)
'''
Training loss: 3.825908660888672
Validation loss: 3.7619335651397705
'''

start_time = time.time()
torch.manual_seed(123)
optimizer = torch.optim.AdamW(
   model.parameters(), lr=0.00005, weight_decay=0.1
)
num_epochs = 2
train_losses, val_losses, tokens_seen = train_model_simple(
   model, train_loader, val_loader, optimizer, device,
   num_epochs=num_epochs, eval_freq=5, eval_iter=5,
   start_context=format_input(val_data[0]), tokenizer=tokenizer
)
end_time = time.time()
execution_time_minutes = (end_time - start_time) / 60
print(f"Training completed in {execution_time_minutes:.2f} minutes.")

file_name = f"{re.sub(r'[ ()]', '', CHOOSE_MODEL) }-sft.pth"
torch.save(model.state_dict(), file_name)
print(f"Model saved as {file_name}")






